from torch.utils.data import random_split, DataLoader, Subset, Dataset
from utils.readers import get_data_dir, get_generator
from torchvision import transforms


class CustomDataset(Dataset):
    def __init__(self, original_dataset, transform=None):
        self.data = original_dataset  # 原始数据

    def __getitem__(self, index):
        x, y = self.data[index]
        return x, y

    def __len__(self):
        return len(self.data)


class DownloadData:
    def __init__(self, batch_size=128):
        self.full = None
        self.train = None
        self.test = None
        self.val = None
        self.data_dir = get_data_dir(__file__)
        self.__test_transform = None
        self.__train_transform = None
        self.batch_size = batch_size

    def prepare_data(self):
        self._download()
        self._divide()

    def _download(self):
        raise NotImplementedError

    def _divide(self):
        if self.full is None or self.test is None:
            raise ValueError("Train, test, and val must be defined")

        train_size = len(self.full) - self.val_size()
        train, val = random_split(
            self.full, [train_size, self.val_size()], generator=get_generator()
        )

        train.dataset.transform = self.train_transform()
        val = Subset(val, val.indices)
        val.dataset.transform = self.test_transform()
        self.test.transform = self.test_transform()
        self.train = CustomDataset(train)
        self.val = CustomDataset(val)
        self.test = CustomDataset(self.test)

    def val_size(self):
        return 5000

    def train_transform(self):
        if self.__train_transform is None:
            self.__train_transform = transforms.ToTensor()
        return self.__train_transform

    def test_transform(self):
        if self.__test_transform is None:
            self.__test_transform = transforms.ToTensor()
        return self.__test_transform


def load_first_batch(dataset, size=100):
    loader = DataLoader(dataset, batch_size=size)
    images, labels = next(iter(loader))
    return images, labels

